#!/usr/bin/env python3
import argparse
import os
import sys
from time import sleep
import struct
import grpc

sys.path.append(
os.path.join(os.path.dirname(os.path.abspath(__file__)),
                 '../../utils/'))
import p4runtime_lib.bmv2
import p4runtime_lib.helper
from p4runtime_lib.error_utils import printGrpcError
from p4runtime_lib.switch import ShutdownAllSwitchConnections

length = 0
def fetch_responses(connection):
    try:
        # 遍历连接的消息响应
        for response in connection.stream_msg_resp:
            if response.WhichOneof("update") == "packet":  # 如果响应是数据包类型
                packet = response.packet
                payload = packet.payload
                l = len(payload)
                h = format(payload[l-3], '02x')  # 将十进制转换为两位十六进制
                if h == "01":  # 如果ECN字段为01
                    print("Receiving response.")
                    print("ECN=1")
                if h == "11":  # 如果ECN字段为11
                    print("Receiving response")
                    print("ECN=3")
                    print("Congestion happens!")  # 发生拥塞
    except grpc.RpcError as e:  # 如果出现 gRPC 错误
        print("Error  happens!")
        pass

def writeForwardRules(p4info_helper, ingress_sw, egress_sw,
                        match_fields, action_params):
    global length  # 全局变量队列长度
    # 构建转发规则条目
    table_entry = p4info_helper.buildTableEntry(
        table_name="MyIngress.ipv4_lpm",
        match_fields={
            "hdr.ipv4.dstAddr": match_fields  # 匹配目的 IPv4 地址
        },
        action_name="MyIngress.ipv4_forward",  # 使用 IPv4 转发动作
        action_params=action_params
    )
    ingress_sw.WriteTableEntry(table_entry)  # 写入转发规则到入口交换机
    
    # 构建设置ECN标记的规则条目
    table_entry = p4info_helper.buildTableEntry(
        table_name="MyEgress.ecn11",
        default_action=True,  # 默认动作
        action_name="MyEgress.ecn_m",  # 使用标记ECN的动作
        action_params={"qlen": length}  # 传递队列长度参数
    )
    ingress_sw.WriteTableEntry(table_entry)  # 写入设置ECN标记规则到入口交换机

def main(p4info_file_path, bmv2_file_path):
    p4info_helper = p4runtime_lib.helper.P4InfoHelper(p4info_file_path)

    try:
        # 创建三个BMv2交换机连接
        s1 = p4runtime_lib.bmv2.Bmv2SwitchConnection(
            name='s1',
            address='127.0.0.1:50051',
            device_id=0,
            proto_dump_file='logs/s1-p4runtime-requests.txt')

        s2 = p4runtime_lib.bmv2.Bmv2SwitchConnection(
            name='s2',
            address='127.0.0.1:50052',
            device_id=1,
            proto_dump_file='logs/s2-p4runtime-requests.txt')

        s3 = p4runtime_lib.bmv2.Bmv2SwitchConnection(
            name='s3',
            address='127.0.0.1:50053',
            device_id=2,
            proto_dump_file='logs/s3-p4runtime-requests.txt')

        # 更新主控权并设置P4程序
        s1.MasterArbitrationUpdate()
        s2.MasterArbitrationUpdate()
        s3.MasterArbitrationUpdate()

        # 在三个交换机上安装P4程序
        s1.SetForwardingPipelineConfig(p4info=p4info_helper.p4info,
                                       bmv2_json_file_path=bmv2_file_path)
        print("Installed P4 Program using SetForwardingPipelineConfig on s1")
        s2.SetForwardingPipelineConfig(p4info=p4info_helper.p4info,
                                       bmv2_json_file_path=bmv2_file_path)
        print("Installed P4 Program using SetForwardingPipelineConfig on s2")
        s3.SetForwardingPipelineConfig(p4info=p4info_helper.p4info,
                                       bmv2_json_file_path=bmv2_file_path)
        print("Installed P4 Program using SetForwardingPipelineConfig on s3")
        
        # 写入转发规则到s1
        writeForwardRules(p4info_helper, s1, s1, 
        ["10.0.1.1", 32], {"dstAddr": "08:00:00:00:01:01", "port": 2})
        writeForwardRules(p4info_helper, s1, s1, 
        ["10.0.1.11", 32], {"dstAddr": "08:00:00:00:01:11", "port": 1})
        writeForwardRules(p4info_helper, s1, s1, 
        ["10.0.2.0", 24], {"dstAddr": "08:00:00:00:02:00", "port": 3})
        writeForwardRules(p4info_helper, s1, s1, 
        ["10.0.3.0", 24], {"dstAddr": "08:00:00:00:03:00", "port": 4})  # 写入更多转发规则
    # 定义网络拓扑和规则
    # 在交换机s2上设置转发规则
    writeForwardRules(p4info_helper, s2, s2, ["10.0.2.2", 32], {"dstAddr": "08:00:00:00:02:02", "port": 2})
    writeForwardRules(p4info_helper, s2, s2, ["10.0.2.22", 32], {"dstAddr": "08:00:00:00:02:22", "port": 1})
    writeForwardRules(p4info_helper, s2, s2, ["10.0.1.0", 24], {"dstAddr": "08:00:00:00:01:00", "port": 3})
    writeForwardRules(p4info_helper, s2, s2, ["10.0.3.0", 24], {"dstAddr": "08:00:00:00:03:00", "port": 4})
    
    # 在交换机s3上设置转发规则
    writeForwardRules(p4info_helper, s3, s3, ["10.0.3.3", 32], {"dstAddr": "08:00:00:00:01:01", "port": 1})
    writeForwardRules(p4info_helper, s3, s3, ["10.0.1.0", 32], {"dstAddr": "08:00:00:00:01:11", "port": 2})
    writeForwardRules(p4info_helper, s3, s3, ["10.0.2.0", 24], {"dstAddr": "08:00:00:00:02:00", "port": 3})

    # 设置克隆会话记录器的入口
    entry = p4info_helper.buildCloneSessionEntry(100, [{"egress_port": 2, "instance": 1}, {"egress_port": 252, "instance": 2}])
    s1.WritePREEntry(entry)
    print("Monitoring network congestion......")
    print("------------------------------------------------")
    # 获取交换机s1的响应信息
    fetch_responses(s1)

# 捕获键盘中断异常
except KeyboardInterrupt:
    print(" Shutting down.")
# 捕获gRPC错误
except grpc.RpcError as e:
    printGrpcError(e)

# 关闭所有交换机连接
ShutdownAllSwitchConnections()

if __name__ == '__main__':
    # 获取输入的ECN长度
    queue_len = int(input("please input the length of ecn: "))
    # 解析命令行参数
    parser = argparse.ArgumentParser(description='P4Runtime Controller')
    parser.add_argument('--p4info', help='p4info proto in text format from p4c',
                        type=str, action="store", required=False,
                        default='./build/ecn.p4.p4info.txt')
    parser.add_argument('--bmv2-json', help='BMv2 JSON file from p4c',
                        type=str, action="store", required=False,
                        default='./build/ecn.json')
    args = parser.parse_args()

    # 检查p4info文件和BMv2 JSON文件是否存在
    if not os.path.exists(args.p4info):
        parser.print_help()
        print("\np4info file not found: %s\nHave you run 'make'?" % args.p4info)
        parser.exit(1)
    if not os.path.exists(args.bmv2_json):
        parser.print_help()
        print("\nBMv2 JSON file not found: %s\nHave you run 'make'?" % args.bmv2_json)
        parser.exit(1)
    # 执行主函数
    main(args.p4info, args.bmv2_json)